from lib import Container, SequenceContainer, encode_bin, decode_bin
from cStringIO import StringIO


#===============================================================================
# exceptions
#===============================================================================
class ConstructError(Exception):
    pass
class FieldError(ConstructError): 
    pass
class EndOfStreamError(FieldError):
    pass
class ValidatorError(ConstructError):
    pass
class AdapterError(ConstructError):
    pass
class RepeaterError(ConstructError):
    pass
class SwitchError(ConstructError):
    pass
class GetSizeError(ConstructError):
    pass


#===============================================================================
# abstract constructs
#===============================================================================
class Construct(object):
    """
    The mother of all constructs; defines all API methods. Abstract class.
    
    Note: all constructs must have a name. This `name` can be either a string,
    in which case it's expected to be a valid python identifier, or None. All
    other values will be rejected (this is done to aid debugging)
    """
    __slots__ = ["name"]
    def __init__(self, name):
        assert name is None or type(name) == str
        self.name = name
    def __repr__(self):
        return "<%s(%r)>" % (self.__class__.__name__, self.name)
    
    def getsize(self, context = None):
        """returns the size of this construct, based on the context; note that
        not all constructs can be queried for size"""
        if context is None:
            context = {}
        return self._getsize(context)
    
    def parse_stream(self, stream):
        """parses the given stream (file, StringIO, ...). The stream must 
        support read()ing; for some constructs, it must also be seek()able"""
        return self._parse(stream, {})
    def parse(self, buf):
        """parses the given buffer (a string or a buffer object). This method 
        simply creates a stream object and passes it to parse_stream"""
        return self.parse_stream(StringIO(buf))
    def parse_iter(self, stream):
        """iterates over the stream, parsing a single construct every time.
        iteration will stop when an EndOfStreamError occurs"""
        try:
            while True:
                yield self._parse(stream, {})
        except EndOfStreamError:
            pass
    
    def build_stream(self, obj, stream):
        """builds the given object into the stream (file, StringIO, ...). The 
        stream must support write()ing; for some constructs, it must also be
        be seek()able"""
        self._build(obj, stream, {})
    def build(self, obj):
        """builds the given object; returns a string. This method simply creates
        a stream object for build_stream and returns the stream buffer"""
        buf = StringIO()
        self.build_stream(obj, buf)
        return buf.getvalue()
    def build_iter(self, stream, seq):
        """builds a sequence of objects into the stream, one after the other"""
        for obj in seq:
            self._build(obj, stream, {})
    
    def _getsize(self, context):
        raise NotImplementedError
    def _parse(self, stream, context):
        raise NotImplementedError
    def _build(self, obj, stream, context):
        raise NotImplementedError

class Subconstruct(Construct):
    """
    Subconstructs are constructs that wrap another construct. Abstract class.
    """
    __slots__ = ["subcon"]
    def __init__(self, subcon):
        Construct.__init__(self, subcon.name)
        self.subcon = subcon
    def _build(self, obj, stream, context):
        self.subcon._build(obj, stream, context)
    def _parse(self, stream, context):
        return self.subcon._parse(stream, context)
    def _getsize(self, context):
        return self.subcon._getsize(context)

class Adapter(Subconstruct):
    """
    Adapters convert one object into another. Abstract class.
    When parsing, the 'primitive' object is decoded into a more meaningful, 
    'high level' one; when building, the 'high level' object is encoded back 
    into its 'primitive' form. Adapters only change the 'representation' of 
    an object, and are thus expected to be symmetrical: 
    decode(encode(x)) == x; encode(decode(x)) == x
    """ 
    __slots__ = []
    def _parse(self, stream, context):
        return self._decode(self.subcon._parse(stream, context))
    def _build(self, obj, stream, context):
        self.subcon._build(self._encode(obj), stream, context)
    def _decode(self, obj):
        return obj
    def _encode(self, obj):
        return obj

class Validator(Subconstruct):
    """
    Validators validate ('make sure') that the object being parsed or built
    meets a certain condition; if it does not, an exception is raised. 
    Abstract class.
    """
    __slots__ = []
    def _parse(self, stream, context):
        obj = self.subcon._parse(stream, context)
        if not self._validate(obj):
            raise ValidatorError("invalid object", obj)
        return obj
    def _build(self, obj, stream, context):
        if not self._validate(obj):
            raise ValidatorError("invalid object", obj)
        self.subcon._build(obj, stream, context)
    def _validate(self, obj):
        raise NotImplementedError


#===============================================================================
# primitives
#===============================================================================
class Field(Construct):
    """
    A `field` is the most fundamental construct -- it reads or writes a given
    number of raw units. Fields are of a fixed length. See also BitField and
    FormatField.
    
    Constructor:
    * name - the name fo the construct
    * length - the length (in units) to read/write
    
    Example:
    Field("blah", 5)
    """
    __slots__ = ["length"]
    def __init__(self, name, length):
        assert length > 0
        Construct.__init__(self, name)
        self.length = length
    def _parse(self, stream, context):
        data = stream.read(self.length)
        if len(data) < self.length:
            raise EndOfStreamError
        return data
    def _build(self, obj, stream, context):
        if len(obj) != self.length:
            raise FieldError("incorrect length; expected %d, found %d" % 
                (self.length, len(obj)))
        stream.write(obj)
    def _getsize(self, context):
        return self.length

class MetaField(Construct):
    """
    A field of a computed length
    
    Constructor:
    * name - the name of the construct
    * lengthfunc - a function that takes the context as an argument and returns
      the length of the field
    
    Example:
    Struct("blah",
        UBInt8("count"),
        MetaField("data", lambda ctx: ctx["count"])
    )
    """
    __slots__ = ["lengthfunc"]
    def __init__(self, name, lengthfunc):
        Construct.__init__(self, name)
        self.lengthfunc = lengthfunc
    def _parse(self, stream, context):
        length = self.lengthfunc(context.get("_", {}))
        data = stream.read(length)
        if len(data) < length:
            raise EndOfStreamError
        return data
    def _build(self, obj, stream, context):
        length = self.lengthfunc(context.get("_", {}))
        if len(obj) != length:
            raise FieldError("incorrect length; expected %d, found %d" % 
                (length, len(obj)))
        stream.write(obj)
    def _getsize(self, context):
        return self.lengthfunc(context.get("_", {}))

class Struct(Construct):
    """
    Defines a fixed sequence of subconstructs; they are parsed/built exactly in
    that order, without any alignment. If padding/alignment is needed, you will 
    need to add it yourself.
    
    Constructor:
    * name - the name of the construct
    * subcons - the sequence of subconstruts
    
    Example:
    Struct("blah",
        UBInt8("a"),
        ULInt16("b"),
        UBInt32("c"),
    )
    """
    __slots__ = ["subcons"]
    def __init__(self, name, *subcons):
        Construct.__init__(self, name)
        self.subcons = subcons
    
    @staticmethod
    def _parse_sc(sc, obj, stream, context):
        if sc.name is None:
            sc._parse(stream, {"_" : context})
        elif sc.name == "<<EMBED>>":
            subobj = sc._parse(stream, context)
            obj.__dict__.update(subobj.__dict__)
        else:
            subobj = sc._parse(stream, {"_" : context})
            context[sc.name] = subobj
            setattr(obj, sc.name, subobj)
    @staticmethod
    def _build_sc(sc, obj, stream, context):
        if sc.name is None:
            sc._build(None, stream, {None : context})
        elif sc.name == "<<EMBED>>":
            sc._build(obj, stream, context)
        else:
            subobj = getattr(obj, sc.name)
            sc._build(subobj, stream, {"_" : context})
            context[sc.name] = subobj
    
    def _parse(self, stream, context):
        obj = Container()
        for sc in self.subcons:
            self._parse_sc(sc, obj, stream, context)
        return obj
    def _build(self, obj, stream, context):
        for sc in self.subcons:
            self._build_sc(sc, obj, stream, context)
    def _getsize(self, context):
        total = 0
        for sc in self.subcons:
            if sc.name is None:
                subcontext = None
            elif sc.name == "<<EMBED>>":
                subcontext = context
            else:
                subcontext = context.get(sc.name, {})
            total += sc._getsize(subcontext)
        return total

class BitStruct(Struct):
    """
    Like Struct, but works at the bit-level. When parsing, it first reads its
    size from the stream, converts it to a bit string (\\x00, \\x01), and only 
    then parses the subconstruts. When building, it builds to a separate 
    bit-string stream, converts it to bytes, and only then writes it to the
    actual stream.
    Note: due to these constraints, the total size of the subconstruts must be 
    a multiple of 8.

    Constructor:
    * name - the name of the construct
    * subcons - the sequence of subconstruts
    
    Example:
    BitStruct("blah",
        Nibble("a"), # 4 bits
        Flag("b"),   # 1 bit
        Padding(3),  # 3 bits
    )
    """
    __slots__ = []
    def _parse(self, stream, context):
        data = stream.read(self._getsize(context))
        bitstream = StringIO(encode_bin(data))
        return Struct._parse(self, bitstream, context)
    def _build(self, obj, stream, context):
        bitstream = StringIO()
        Struct._build(self, obj, bitstream, context)
        stream.write(decode_bin(bitstream.getvalue()))
    def _getsize(self, context):
        bitsize = Struct._getsize(self, context)
        assert bitsize & 7 == 0
        return bitsize // 8

class Union(Struct):
    """
    Like Struct, but does not advance the stream position after each 
    subconstrut. Allows multiple 'views' of the same data (i.e., treat it once 
    like an integer, then like a floating point); very similar to unions in C. 
    When parsing, it will restore the stream position after each subconstrut
    has parsed. When all have finished, the stream's position is set to the
    maximal position reached.
    When building, all subconstruts are built, but ONLY THE FIRST ONE is 
    written to the stream, therefore you must take care over which subcon is
    defined first.
    
    Constructor:
    * name - the name of the construct
    * subcons - the sequence of subconstruts
    
    Example:
    Union("blah",
        UBInt16("one_word"),
        Struct("two_bytes",
            UBInt16("high"),
            UBInt16("low"),
        )
    )
    """
    __slots__ = ["first_subcon", "subcons"]
    def __init__(self, name, *subcons):
        assert len(subcons) > 0, "expected at least one subconstruct"
        Struct.__init__(self, name, *subcons[1:])
        first_subcon = subcons[0]
    def _parse(self, stream, context):
        obj = Container()
        startpos = stream.tell()
        # first subcon
        self._parse_sc(self.first_subcon, obj, stream, context)
        finalpos = stream.tell()
        # rest subcons
        for sc in self.subcons:
            stream.seek(startpos)
            self._parse_sc(sc, obj, stream, context)
        stream.seek(finalpos)
        return obj
    def _build(self, obj, stream, context):
        self._build_sc(self.first_subcon, obj, stream, context)
        dummy_stream = StringIO()
        for sc in self.subcons:
            self._build_sc(sc, obj, dummy_stream, context)
            dummy_stream.reset()
    def _getsize(self, context):
        return self.subcons[0]._getsize(context)

class Optional(Subconstruct):
    """
    Optional subconstruct: when parsing, if it fails to parse, a default value 
    will be returned instead. When building, if the object is the default value,
    it will be skipped.
    
    Constructor:
    * subcon - the subconstruct to optionally parse/build
    * default - the default value; by default it is None
    
    Example:
    Optional(UBInt8("blah"))
    """
    __slots__ = ["default"]
    def __init__(self, subcon, default = None):
        Subconstruct.__init__(self, subcon)
        self.default = default
    def _parse(self, stream, context):
        try:
            return self.subcon._parse(stream, context)
        except ConstructError:
            return self.default
    def _build(self, obj, stream, context):
        if obj != self.default:
            self.subcon._build(obj, stream, context)

class Repeater(Subconstruct):
    """
    Repeats the given subconstruct for at least `min` times, up to `max` times.
    It will attempt to read `max` times, unless it fails. In any way, it will
    require that at least `min` successful repetitions have made, or it will
    raise an exception.
    
    Constructor:
    * min - the minimal number of times to repeat; `min` must be >= 0
    * max - the maximal number of times to repeat; `max` must be >= `min`
    * subcon - the subconstruct to repeat
    
    Notes:
    * getsize() will work only if `min` == `max`
    
    Example:
    Repeater(UBInt8("blah"), 10, 20)
    """
    __slots__ = ["min", "max"]
    def __init__(self, min, max, subcon):
        Subconstruct.__init__(self, subcon)
        assert min >= 0
        assert max >= min
        self.min = min
        self.max = max
    def _parse(self, stream, context):
        cnt = 0
        obj = SequenceContainer()
        try:
            while cnt < self.max:
                pos = stream.tell()
                obj.append(self.subcon._parse(stream, context))
                cnt += 1
        except ConstructError:
            stream.seek(pos) 
        if cnt < self.min:
            raise RepeaterError("expected %d..%d, found %d" % 
                (self.min, self.max, cnt))
        return obj
    def _build(self, obj, stream, context):
        cnt = len(obj)
        if cnt < self.min or cnt > self.max:
            raise RepeaterError("expected %d..%d, found %d" % 
                (self.min, self.max, cnt))
        for subobj in obj:
            self.subcon._build(subobj, stream, context)
    def _getsize(self, context):
        if self.min != self.max:
            raise GetSizeError("can't precalculate the size (min != max)")
        return self.min * self.subcon._getsize(context)

class RepeatUntil(Subconstruct):
    """
    Repeats the given subconstruct until the predicate function indicates to
    stop (i.e. repeat until some condition is met)
    
    Constructor:
    * pedicate - a predicate function (returns True or False) that accepts
      (obj, context) as its arguments; when the function returns True, the 
      repeater finishes.
    * subcon - the sub construct to repeat
    
    Notes: 
    * getsize() isn't supported for this construct
    
    Example:
    RepeatUntil(lambda obj, ctx: obj != 0,
        UBInt8("blah")
    )
    """
    __slots__ = ["predicate"]
    def __init__(self, predicate, subcon):
        Subconstruct.__init__(self, subcon)
        self.predicate = predicate
    def _parse(self, stream, context):
        obj = SequenceContainer()
        subctx = context.get("_", {})
        while True:
            subobj = self.subcon._parse(stream, context)
            obj.append(subobj)
            if self.predicate(subobj, subctx):
                break
        return obj
    def _build(self, obj, stream, context):
        subctx = context.get("_", {})
        for subobj in obj:
            self.subcon._build(subobj, stream, context)
            if self.predicate(subobj, subctx):
                break
        else:
            raise RepeaterError("unterminated sequence")
    def _getsize(self, context):
        raise GetSizeError("can't precalculate the size")

class MetaRepeater(Subconstruct):
    """
    Repeats the given subconstruct for a computed number of times. 
    
    Constructor:
    * countfunc - a function that takes the context as an argument and returns
      the number of repetitions to perform
    * subcon - the subconstruct to repeat
    
    Example:
    Struct("blah",
        UBInt8("count"),
        MetaRepeater(lambda ctx: ctx["count"],
            UBInt8("data")
        )
    )
    """
    __slots__ = ["countfunc"]
    def __init__(self, countfunc, subcon):
        Subconstruct.__init__(self, subcon)
        self.countfunc = countfunc
    def _parse(self, stream, context):
        cnt = self.countfunc(context.get("_", {}))
        obj = SequenceContainer([None] * cnt)
        i = 0
        while i < cnt:
            obj[i] = self.subcon._parse(stream, context)
            i += 1
        return obj
    def _build(self, obj, stream, context):
        cnt = self.countfunc(context.get("_", {}))
        if cnt != len(obj):
            raise RepeaterError("expected %d, found %d" % (cnt, len(obj)))
        for subobj in obj:
            self.subcon._build(subobj, stream, context)
    def _getsize(self, context):
        cnt = self.countfunc(context.get("_", {}))
        return cnt * self.subcon._getsize(context)

class Switch(Construct):
    """
    Selects the case to use from the `cases` dict using the `selector` function. 
    
    Constructor:
    * name - the name of the construct
    * selector - a function that takes the context as an argument and returns 
      the key to use
    * cases - a dictionary of (key, construct) items, from which the selector
      function selects the case
    * default - the default construct to use, in case the key does not exist in
      the cases dict; if no default case is given, an error is raised.
    
    Example:
    Struct("blah",
        UBInt8("type"),
        Switch("value", lambda ctx: ctx["type"],
            {
                1 : UBInt8("data"),
                2 : UBInt16("data"),
                3 : UBInt32("data"),
            }
        )
    """
    __slots__ = ["selector", "cases", "default"]
    
    class nodefault(Construct):
        __slots__ = []
        def _error(self, *args):
            raise SwitchError("no default case")
        _parse = _build = _getsize = _error
    nodefault = nodefault(None)
    
    def __init__(self, name, selector, cases, default = nodefault):
        Construct.__init__(self, name)
        self.selector = selector
        self.cases = cases
        self.default = default
    def _parse(self, stream, context):
        key = self.selector(context.get("_", {}))
        case = self.cases.get(key, self.default)
        return case._parse(stream, context)
    def _build(self, obj, stream, context):
        key = self.selector(context.get("_", {}))
        case = self.cases.get(key, self.default)
        case._build(obj, stream, context)
    def _getsize(self, context):
        key = self.selector(context.get("_", {}))
        case = self.cases.get(key, self.default)
        return case._getsize(context)

class Pointer(Subconstruct):
    """
    Diverts parsing/building to a different stream position, pointed to by
    `offsetfunc`. Useful for data structures that use pointers, such as the 
    tables of the ELF headers, etc.
    
    Constructor:
    * offsetfunc - a function that takes the context as an argument and 
      returns the absolute stream position to jump to. 
    * subcon - the subconstruct that will be parsed/built at that position
    
    Notes: 
    * in order to use relative positions, just place an Achor somewhere, and 
      add it's value to the relative offset.
    * it would make sense to combine pointers and OnDemand constructs.
    
    Example:
    Struct("blah",
        Anchor("begining"),
        UBInt32("offset"),
        Pointer(lambda ctx: ctx["beginning"] + ctx["offset"],
            UBInt32("pointed_data")
        )
    )
    """
    __slots__ = ["offsetfunc"]
    def __init__(self, offsetfunc, subcon):
        Subconstruct.__init__(self, subcon)
        self.offsetfunc = offsetfunc
    def _parse(self, stream, context):
        oldpos = stream.tell()
        newpos = self.offsetfunc(context["_"])
        stream.seek(newpos)
        obj = self.subcon._parse(stream, context)
        stream.seek(oldpos)
        return obj
    def _build(self, obj, stream, context):
        oldpos = stream.tell()
        newpos = self.offsetfunc(context["_"])
        stream.seek(newpos)
        self.subcon._build(obj, stream, context)
        stream.seek(oldpos)

class LazyBound(Construct):
    """
    A lazy-bound construct; allows the definition of recursive constructs (like
    linked lists)
    
    Constructor:
    * name - the name of this construct
    * lazyfunc - a function that takes no arguments and returns the construct
      to use
    
    Example:
    node = Struct("blah",
        UBInt8("value"),
        LazyBound("next", lambda : node) # warning: endless recursion
    )
    """
    __slots__ = ["lazyfunc"]
    def __init__(self, name, lazyfunc):
        Construct.__init__(self, name)
        self.lazyfunc = lazyfunc
    def _parse(self, stream, context):
        return self.lazyfunc()._parse(stream, context)
    def _build(self, obj, stream, context):
        return self.lazyfunc()._build(obj, stream, context)
    def _getsize(self, context):
        return self.lazyfunc()._getsize(context)

class OnDemand(Subconstruct):
    """
    On-demand parsing/building of constructs. When parsing, it will return an
    OnDemandContainer object, which will actually parse the object only upon
    request (by accessing it's '.value' property).
    When building, the object will be built only if 'build_always' is set to
    True, or if it has a value (the value has been 'demanded'); otherwise, 
    it will just be skipped, advancing the stream by its size.
    Note that only OnDemandContainer objects can be built. You can create an
    OnDemandContainer object manually by wrapping the 'value' by it:
    odcobj = OnDemandContainer(value)
    
    Constructor:
    * subcon - the subconstruct to parse/build on demand
    * build_always - a boolean indicating whether the subconstruct should be
      demanded prior to building it. the default is False, which means the 
      subconstruct will be skipped unless it was demanded externally.
      
    Example:
    # this big string will be actually read only on demand
    OnDemand(String("blah", 10000)) 
    """
    __slots__ = ["build_always"]
    def __init__(self, subcon, build_always = False):
        Subconstruct.__init__(self, subcon)
        self.build_always = build_always
    def _parse(self, stream, context):
        obj = OnDemandContainer.create_lazy(self.subcon, stream, context)
        stream.seek(obj.size, 1)
        return obj
    def _build(self, obj, stream, context):
        if self.build_always or obj.has_value:
            self.subcon._build(obj.value, stream, context)
        else:
            stream.seek(self.subcon._getsize(context), 1)



